import random
import numpy as np

from beam.util_psll import get_psll_by_phase

from util.util_ris_pattern import point_2_phase, eps
from util.util_log import setup_logging

from util.util_image import draw_img_xyz_2, draw_img

#======================================== LOSS (adapter MSE & SSIM & NCC) ===========================================
import math

from util.util_ris_pattern import phase_2_bit, phaseBit_2_deg, phase_2_pattern_xyz_fft, phase_2_pattern
from skimage.metrics import structural_similarity as ssim

def get_patterndbw(phaseRad, bit_num=1):
    # 计算方向图
    # patternBit = phase_2_pattern(phaseRad)                # 公式法直接计算, 准确但速度太慢
    patternBit, x, y = phase_2_pattern_xyz_fft(phaseRad)  # FFT法计算, 快速
    # pattern 转dB
    pattern_dbw = 20 * np.log10(np.abs(patternBit) / np.max(np.max(np.abs(patternBit))) + eps)
    return pattern_dbw


def get_ncc(pattern_true, pattern_pred):
    # 计算均值
    mean_true = np.mean(pattern_true)
    mean_pred = np.mean(pattern_pred)
    # 计算标准差
    std_true = np.std(pattern_true)
    std_pred = np.std(pattern_pred)
    # 计算分子和分母
    numerator = np.sum((pattern_true - mean_true) * (pattern_pred - mean_pred))
    denominator = (pattern_true.size - 1) * std_true * std_pred
    # 计算NCC
    ncc = numerator / denominator if denominator != 0 else 0
    return ncc

def get_mse(pattern_true, pattern_pred, threshold=100):
    mse = np.mean((pattern_true - pattern_pred) ** 2)
    # 假设 pattern_true 和 pattern_pred 的取值范围在 [0, -threshold] 之间
    max_mse = threshold ** 2
    normalized_mse = 1 - (mse / max_mse)
    # 确保 normalized_mse 在 [0, 1] 范围内
    normalized_mse = np.clip(normalized_mse, 0, 1)
    return normalized_mse

def get_ssim(pattern_true, pattern_pred):
    ssim_index = ssim(pattern_true, pattern_pred, data_range=pattern_true.max() - pattern_true.min())
    return ssim_index


def loss_adapter_mse_ssim_ncc(pattern_true, phase_pred):
    """
    自适应alpha, alpha * pattern-SSIM + (1-alpha) * pattern-MSE
    """
    # 计算pred方向图
    pattern_pred = get_patterndbw(phase_pred)
    # draw_img(pattern_pred)
    # draw_img(pattern_true)
    # 计算方向图的MSE和SSIM
    # pattern_mse = get_mse(pattern_true, pattern_pred)
    # pattern_ssim = get_ssim(pattern_true, pattern_pred)
    pattern_ncc = get_ncc(pattern_true, pattern_pred) * 1000
    return pattern_ncc
    # 倒S型参数 y=cos(90*ssim), 修改为张量计算
    # s_para = math.cos(math.radians(90 * pattern_ssim))
    # 自适应权重加权组合
    # total_loss = s_para * pattern_ssim + (1 - s_para) * pattern_mse
    # logger.info(f"pattern_mse={pattern_mse}, pattern_ssim={pattern_ssim}, pattern_ncc={pattern_ncc}, s_para={s_para}, total_loss={total_loss}")
    # return total_loss


#===================================================== mask ====================================================
def loss_mask_mse(B_add, B1, B2, C1, C2, alpha, beta):
    # 计算峰值保留损失
    peak_loss = np.sum(np.abs((B_add * C1) - (B1 * C1)) ** 2) + \
                np.sum(np.abs((B_add * C2) - (B2 * C2)) ** 2)
    # 计算背景抑制损失
    background_loss = np.sum(np.abs(B_add * (1 - C1)) ** 2) + \
                      np.sum(np.abs(B_add * (1 - C2)) ** 2)
    # 计算总损失
    total_loss = alpha * peak_loss + beta * background_loss
    print(f"total_loss={total_loss}, peak_loss={peak_loss}, background_loss={background_loss}")
    return total_loss

def loss_mask_mse_2(B_add, C1, C2, alpha, beta):
    # 计算峰值保留损失
    peak_loss = np.sum(np.abs(B_add * C1) ** 2) + np.sum(np.abs(B_add * C2) ** 2)
    # 计算背景抑制损失
    background_loss = np.sum(np.abs(B_add * (1 - C1)) ** 2) + np.sum(np.abs(B_add * (1 - C2)) ** 2)
    # 计算总损失
    total_loss = alpha * peak_loss + beta * background_loss
    print(f"total_loss={total_loss}, peak_loss={peak_loss}, background_loss={background_loss}")
    return total_loss


from skimage.metrics import structural_similarity as ssim

def loss_mask_ssim(B_add, B1, B2, C1, C2, alpha, beta):
    # 计算峰值保留部分的SSIM
    peak_ssim_1 = ssim((B_add * C1), (B1 * C1), data_range=B_add.max() - B_add.min())
    peak_ssim_2 = ssim((B_add * C2), (B2 * C2), data_range=B_add.max() - B_add.min())
    peak_loss = 1 - ((peak_ssim_1 + peak_ssim_2) / 2)

    # 计算背景抑制部分的SSIM
    background_ssim_1 = ssim((B_add * (1 - C1)), np.zeros_like(B_add * (1 - C1)), data_range=B_add.max() - B_add.min())
    background_ssim_2 = ssim((B_add * (1 - C2)), np.zeros_like(B_add * (1 - C2)), data_range=B_add.max() - B_add.min())
    background_loss = 1 - ((background_ssim_1 + background_ssim_2) / 2)

    # 计算总损失
    total_loss = alpha * peak_loss + beta * background_loss
    print(f"ssim: total_loss={total_loss}, peak_loss={peak_loss}, background_loss={background_loss}")
    return total_loss


#===================================================== GA ======================================================
def get_patterndbw_xy(phaseRad):
    # 计算方向图
    # patternBit = phase_2_pattern(phaseRad)                # 公式法直接计算, 准确但速度太慢
    patternBit, x, y = phase_2_pattern_xyz_fft(phaseRad)  # FFT法计算, 快速
    # pattern 转dB
    pattern_dbw = 20 * np.log10(np.abs(patternBit) / np.max(np.max(np.abs(patternBit))) + eps)
    return pattern_dbw, x, y


from util.util_analysis_plane import get_peaks, get_peak_nth

def get_psll(pattern_dbw, beam_num):
    psll = 0
    # pattern找峰值
    peaks = get_peaks(pattern_dbw)
    # 找beam+1峰作为 PSLL
    peak_2nd = get_peak_nth(peaks, beam_num)
    if peak_2nd is not None:
        psll = peak_2nd[0]
    return psll


class GA():
    __bit_num = 0  # 比特数
    __beam_num = 0  # 波束数

    def __init__(self, bit_num, beam_num, population_size=50, num_generations=100, num_parents=10, mutation_rate=0.1):
        # 初始化阵列相关
        self.__bit_num = bit_num
        self.__beam_num = beam_num
        # 初始化遗传算法相关
        self.population_size = population_size
        self.num_generations = num_generations
        self.num_parents = num_parents
        self.mutation_rate = mutation_rate
        self.population = None
        self.best_individual = None
        self.best_fitness = None
        self.best_fitness_history = []  # 保存每一代的最佳适应度值
        self.best_individual_history = []  # 保存每一代的最佳适应度值的个体
        # fitness
        self.pattern_target = np.ones((64, 64), dtype=float)

    def initialize_population(self, phase_mix_init):
        """初始化种群"""
        self.population = [phase_mix_init] * self.population_size

    def fitness(self, phase):
        return loss_adapter_mse_ssim_ncc(pattern_true=self.pattern_target, phase_pred=phase)
        # return get_psll_by_phase(phase, self.__bit_num, self.__beam_num)

    def selection(self):
        """选择操作"""
        fitness_scores = [self.fitness(individual) for individual in self.population]
        sorted_indices = np.argsort(fitness_scores)  # 从低到高排序
        selected_parents = [self.population[i] for i in sorted_indices[:self.num_parents]]
        # 获取选中的父代个体对应的适应度值
        selected_fitness_scores = [fitness_scores[i] for i in sorted_indices[:self.num_parents]]
        return selected_parents, selected_fitness_scores

    def crossover(self, parents, offspring_size):
        """交叉操作"""
        offspring = []
        for _ in range(offspring_size):
            parent1, parent2 = random.sample(parents, 2)
            parent1 = np.array(parent1)
            parent2 = np.array(parent2)
            # 获取数组的形状
            shape = parent1.shape
            # 将 parent1 和 parent2 扁平化为一维数组
            flat_parent1 = parent1.flatten()
            flat_parent2 = parent2.flatten()
            # 随机选择一个切割位置
            cut_point = random.randint(0, len(flat_parent1))
            # 生成子代
            child = np.concatenate((flat_parent1[:cut_point], flat_parent2[cut_point:]))
            # 将子代重新塑形为原始的二维数组形状
            child = child.reshape(shape)
            child = child.tolist()
            #
            offspring.append(np.array(child))
        return offspring

    def mutation(self, offspring):
        """变异操作"""
        for individual in offspring:
            if random.random() < self.mutation_rate:
                # 获取数组的数量和长度
                num_clusters, array_length = individual.shape
                # 生成一个与 individual 形状相同的随机数数组
                random_values = np.random.uniform(-180, 180, (num_clusters, array_length))
                # 决定哪些元素需要变异
                mask = np.random.rand(num_clusters, array_length) < self.mutation_rate
                # 应用变异
                individual[mask] = random_values[mask]
        return offspring

    def run(self, phase_mix_init, pattern_target, logger):
        logger.info("population_size=%d, num_generations=%d, num_parents=%d, mutation_rate=%d"
                    % (self.population_size, self.num_generations, self.num_parents, self.mutation_rate))
        # fitness
        self.pattern_target = pattern_target
        # """初始化返回值"""
        self.best_fitness = self.fitness(phase_mix_init)
        self.best_individual = phase_mix_init
        self.best_fitness_history = []
        self.best_individual_history = []
        # """运行遗传算法 -- 初始化阶段"""
        self.initialize_population(phase_mix_init)
        # """运行遗传算法 -- 搜索阶段"""
        for generation in range(self.num_generations):
            # 选择操作
            selected_parents, selected_fitness_scores = self.selection()
            # 交换操作
            offspring = self.crossover(selected_parents, self.population_size - self.num_parents)
            # 变异操作
            offspring = self.mutation(offspring)
            # 计算后代的适应度值
            offspring_fitness_scores = [self.fitness(individual) for individual in offspring]
            # 更新种群
            self.population = selected_parents + offspring
            # 合并父代和后代的适应度值
            all_fitness_scores = selected_fitness_scores + offspring_fitness_scores
            #
            # 找到当前代的最佳个体
            this_best_index = np.argmax(all_fitness_scores)  # 找到最佳适应度值的索引
            if this_best_index < len(selected_parents):
                this_best_individual = selected_parents[this_best_index]  # 最佳个体来自父代
            else:
                this_best_individual = offspring[this_best_index - len(selected_parents)]  # 最佳个体来自后代
            this_best_fitness = all_fitness_scores[this_best_index]  # 最佳适应度值
            #
            # 更新最佳适应度
            if this_best_fitness > self.best_fitness:
                self.best_fitness = this_best_fitness
                self.best_individual = this_best_individual
            # 记录最佳适应度曲线
            self.best_fitness_history.append(self.best_fitness)
            self.best_individual_history.append(self.best_individual)
            #
            # logger.info("generation=%d: self.best_fitness=%f, self.best_individual:%s"
            #             % (generation, self.best_fitness, self.best_individual))
            logger.info("generation=%d: self.best_fitness=%f" % (generation, self.best_fitness))
        return self.best_individual, self.best_fitness, self.best_fitness_history, self.best_individual_history




def main_ga(logger):
    # 测试代码
    bit_num = 1
    theta = 10
    phi = 10
    ga = GA(bit_num=bit_num, beam_num=1, population_size=30, num_generations=6, num_parents=10, mutation_rate=0.1)
    # 1.背投影法(BP)计算码阵
    phaseRad = point_2_phase(theta, phi)
    patternBit_dbw = get_patterndbw(phaseRad, bit_num)
    patternBit_dbw_threld_30 = np.copy(patternBit_dbw)
    # pattern 转dB
    # 将 patternBit_dbw 中所有小于 -3 的值设置为 -100
    patternBit_dbw_threld_30[patternBit_dbw_threld_30 < -3] = -30
    # 2.GA
    best_individual, best_fitness, best_fitness_history, best_individual_history \
        = ga.run(phaseRad, patternBit_dbw_threld_30, logger)
    # 3.计算优化后的pattern
    phaseRad_ga = best_individual
    patternBit_dbw_ga = get_patterndbw(phaseRad_ga, bit_num)
    # 检查PSLL
    psll = get_psll(pattern_dbw=patternBit_dbw,
                    beam_num=1)
    psll_ga = get_psll(pattern_dbw=patternBit_dbw_ga, beam_num=1)
    logger.info(f"psll={psll}, psll_ga={psll_ga}")
    # 画图
    # draw_img(phaseRad)
    # patterndbw, x, y = get_patterndbw_xy(phaseRad)
    # draw_img_xyz_2(patterndbw, x, y)
    # draw_img(phaseRad_ga)
    # patterndbw_ga, x_ga, y_ga = get_patterndbw_xy(phaseRad_ga)
    # draw_img_xyz_2(patterndbw_ga, x_ga, y_ga)


def test_loss_mask_mse():
    # 示例数据
    # B_add = np.array([[1, 2], [3, 4]])
    # B1 = np.array([[0, 2], [3, 0]])
    # B2 = np.array([[1, 0], [0, 4]])
    # C1 = np.array([[1, 0], [0, 1]])
    # C2 = np.array([[0, 1], [1, 0]])
    #
    B_add_1 = np.array([[1, 2, 0, 0], [3, 100, 2, 1], [3, 0, 1, 3], [2, 1, 100, 0]])
    B_add_2 = np.array([[1, 2, 0, 0], [3, 100, 2, 1], [3, 0, 1, 3], [2, 1, 2, 0]])
    B_add_3 = np.array([[1, 2, 0, 0], [3, 4, 2, 1], [3, 0, 1, 3], [2, 1, 100, 0]])
    B_add_4 = np.array([[1, 2, 0, 0], [3, 4, 2, 1], [3, 0, 1, 3], [2, 1, 2, 0]])
    B_add_5 = np.array([[1, 2, 0, 0], [3, 4, 2, 1], [3, 0, 100, 3], [2, 1, 2, 0]])
    B_add_6 = np.array([[1, 2, 0, 0], [3, 100, 2, 1], [3, 0, 100, 3], [2, 1, 2, 0]])
    B_add_7 = np.array([[1, 2, 0, 0], [3, 4, 2, 1], [3, 0, 100, 3], [2, 1, 100, 0]])
    #
    B1 = np.array([[0, 2, 1, 3], [3, 100, 0, 2], [1, 1, 0, 1], [2, 1, 3, 1]])
    B2 = np.array([[1, 0, 1, 2], [0, 4, 0, 2], [1, 2, 2, 3], [0, 2, 100, 2]])
    C1 = np.array([[0, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]])
    C2 = np.array([[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 1, 0]])
    #
    alpha = 1
    beta = 0.1
    # 计算损失
    loss_case_1 = loss_mask_mse(B_add_1, B1, B2, C1, C2, alpha, beta)
    print(f"loss_case_1: {loss_case_1}")
    loss_case_2 = loss_mask_mse(B_add_2, B1, B2, C1, C2, alpha, beta)
    print(f"loss_case_2: {loss_case_2}")
    loss_case_3 = loss_mask_mse(B_add_3, B1, B2, C1, C2, alpha, beta)
    print(f"loss_case_3: {loss_case_3}")
    loss_case_4 = loss_mask_mse(B_add_4, B1, B2, C1, C2, alpha, beta)
    print(f"loss_case_4: {loss_case_4}")
    loss_case_5 = loss_mask_mse(B_add_5, B1, B2, C1, C2, alpha, beta)
    print(f"loss_case_5: {loss_case_5}")
    loss_case_6 = loss_mask_mse(B_add_6, B1, B2, C1, C2, alpha, beta)
    print(f"loss_case_6: {loss_case_6}")
    loss_case_7 = loss_mask_mse(B_add_7, B1, B2, C1, C2, alpha, beta)
    print(f"loss_case_7: {loss_case_7}")

from multi_beam.multi_beam_PS_complex import psm_complex_beam_n

def test_pattern_loss_mask_mse():
    # 示例数据
    bit_num = 1
    # 1.背投影法(BP)计算码阵
    phaseRad_1 = point_2_phase(30, 0)
    phaseRad_2 = point_2_phase(30, 80)
    phaseRad_3 = point_2_phase(30, 150)
    phaseRad_4 = point_2_phase(30, 300)
    # 2.
    patternBit_dbw_1 = get_patterndbw(phaseRad_1, bit_num)
    patternBit_dbw_2 = get_patterndbw(phaseRad_2, bit_num)
    patternBit_dbw_3 = get_patterndbw(phaseRad_3, bit_num)
    patternBit_dbw_4 = get_patterndbw(phaseRad_4, bit_num)
    patternBit_dbw_1_mask = np.copy(patternBit_dbw_1)
    patternBit_dbw_2_mask = np.copy(patternBit_dbw_2)
    patternBit_dbw_3_mask = np.copy(patternBit_dbw_3)
    patternBit_dbw_4_mask = np.copy(patternBit_dbw_4)
    threshold = -3
    patternBit_dbw_1_mask[patternBit_dbw_1_mask >= threshold] = 1
    patternBit_dbw_1_mask[patternBit_dbw_1_mask < threshold] = 0
    patternBit_dbw_2_mask[patternBit_dbw_2_mask >= threshold] = 1
    patternBit_dbw_2_mask[patternBit_dbw_2_mask < threshold] = 0
    patternBit_dbw_3_mask[patternBit_dbw_3_mask >= threshold] = 1
    patternBit_dbw_3_mask[patternBit_dbw_3_mask < threshold] = 0
    patternBit_dbw_4_mask[patternBit_dbw_4_mask >= threshold] = 1
    patternBit_dbw_4_mask[patternBit_dbw_4_mask < threshold] = 0
    print(f"mask-1: count(1)={(patternBit_dbw_1_mask == 1).sum()}, count(0)={(patternBit_dbw_1_mask == 0).sum()}")
    print(f"mask-2: count(1)={(patternBit_dbw_2_mask == 1).sum()}, count(0)={(patternBit_dbw_2_mask == 0).sum()}")
    print(f"mask-3: count(1)={(patternBit_dbw_3_mask == 1).sum()}, count(0)={(patternBit_dbw_3_mask == 0).sum()}")
    print(f"mask-4: count(1)={(patternBit_dbw_4_mask == 1).sum()}, count(0)={(patternBit_dbw_4_mask == 0).sum()}")
    # 3.
    phase_mix_12, phaseBit_mix_12, phaseBitDeg_mix_12 \
        = psm_complex_beam_n(phases=[phaseRad_1, phaseRad_2], bit_num=bit_num)
    patternBit_dbw_12 = get_patterndbw(phase_mix_12, bit_num)
    phase_mix_13, phaseBit_mix_13, phaseBitDeg_mix_13 \
        = psm_complex_beam_n(phases=[phaseRad_1, phaseRad_3], bit_num=bit_num)
    patternBit_dbw_13 = get_patterndbw(phase_mix_13, bit_num)
    phase_mix_23, phaseBit_mix_23, phaseBitDeg_mix_23 \
        = psm_complex_beam_n(phases=[phaseRad_2, phaseRad_3], bit_num=bit_num)
    patternBit_dbw_23 = get_patterndbw(phase_mix_23, bit_num)
    phase_mix_34, phaseBit_mix_34, phaseBitDeg_mix_34 \
        = psm_complex_beam_n(phases=[phaseRad_3, phaseRad_4], bit_num=bit_num)
    patternBit_dbw_34 = get_patterndbw(phase_mix_34, bit_num)
    #
    alpha, beta = 30000, 1
    # 计算损失
    # loss_case_1_good_2 = loss_mask_mse(patternBit_dbw_12,
    #                                    patternBit_dbw_1, patternBit_dbw_2,
    #                                    patternBit_dbw_1_mask, patternBit_dbw_2_mask,
    #                                    alpha, beta)
    # print(f"loss_case_1_good_2: {loss_case_1_good_2}")
    # loss_case_2_good_1_bad_1_13 = loss_mask_mse(patternBit_dbw_13,
    #                                            patternBit_dbw_1, patternBit_dbw_2,
    #                                            patternBit_dbw_1_mask, patternBit_dbw_2_mask,
    #                                            alpha, beta)
    # print(f"loss_case_2_good_1_bad_1_13: {loss_case_2_good_1_bad_1_13}")
    # loss_case_3_good_1_bad_1_23 = loss_mask_mse(patternBit_dbw_23,
    #                                             patternBit_dbw_1, patternBit_dbw_2,
    #                                             patternBit_dbw_1_mask, patternBit_dbw_2_mask,
    #                                             alpha, beta)
    # print(f"loss_case_3_good_1_bad_1_23: {loss_case_3_good_1_bad_1_23}")
    # loss_case_4_bad_2 = loss_mask_mse(patternBit_dbw_34,
    #                                   patternBit_dbw_1, patternBit_dbw_2,
    #                                   patternBit_dbw_1_mask, patternBit_dbw_2_mask,
    #                                   alpha, beta)
    # print(f"loss_case_4_bad_2: {loss_case_4_bad_2}")
    #
    loss_case_1_good_2 = loss_mask_ssim(patternBit_dbw_12,
                                       patternBit_dbw_1, patternBit_dbw_2,
                                       patternBit_dbw_1_mask, patternBit_dbw_2_mask,
                                       alpha, beta)
    print(f"loss_case_1_good_2: {loss_case_1_good_2}")
    loss_case_2_good_1_bad_1_13 = loss_mask_ssim(patternBit_dbw_13,
                                                patternBit_dbw_1, patternBit_dbw_2,
                                                patternBit_dbw_1_mask, patternBit_dbw_2_mask,
                                                alpha, beta)
    print(f"loss_case_2_good_1_bad_1_13: {loss_case_2_good_1_bad_1_13}")
    loss_case_3_good_1_bad_1_23 = loss_mask_ssim(patternBit_dbw_23,
                                                patternBit_dbw_1, patternBit_dbw_2,
                                                patternBit_dbw_1_mask, patternBit_dbw_2_mask,
                                                alpha, beta)
    print(f"loss_case_3_good_1_bad_1_23: {loss_case_3_good_1_bad_1_23}")
    loss_case_4_bad_2 = loss_mask_ssim(patternBit_dbw_34,
                                      patternBit_dbw_1, patternBit_dbw_2,
                                      patternBit_dbw_1_mask, patternBit_dbw_2_mask,
                                      alpha, beta)
    print(f"loss_case_4_bad_2: {loss_case_4_bad_2}")


import numpy as np
from collections import deque


def find_peak_and_mark(patternBit_dbw_mask):
    # 获取数组的形状
    rows, cols = patternBit_dbw_mask.shape

    # 找到最大值及其索引
    max_value = np.max(patternBit_dbw_mask)
    max_index = np.unravel_index(np.argmax(patternBit_dbw_mask), patternBit_dbw_mask.shape)

    # 初始化结果数组
    result = np.zeros_like(patternBit_dbw_mask, dtype=int)
    result[max_index] = 1

    # 使用队列进行广度优先搜索
    queue = deque([max_index])
    directions = [(-1, 0), (1, 0), (0, -1), (0, 1)]

    while queue:
        current_row, current_col = queue.popleft()

        for dr, dc in directions:
            new_row, new_col = current_row + dr, current_col + dc

            if 0 <= new_row < rows and 0 <= new_col < cols:
                if patternBit_dbw_mask[new_row, new_col] < patternBit_dbw_mask[current_row, current_col]:
                    result[new_row, new_col] = 1
                    queue.append((new_row, new_col))

    return result


def get_mask_by_slope(bit_num, phaseRad):
    patternBit_dbw = get_patterndbw(phaseRad, bit_num)
    patternBit_dbw_mask = np.copy(patternBit_dbw)
    print("ggggggggggggggggggggggggggggggggggggggggggggggggggggggggg")
    draw_img(patternBit_dbw_mask)
    patternBit_dbw_mask = find_peak_and_mark(patternBit_dbw_mask)
    return patternBit_dbw, patternBit_dbw_mask


def get_mask_by_threshold(bit_num, threshold, phaseRad):
    patternBit_dbw = get_patterndbw(phaseRad, bit_num)
    patternBit_dbw_mask = np.copy(patternBit_dbw)
    patternBit_dbw_mask[patternBit_dbw_mask >= threshold] = 1
    patternBit_dbw_mask[patternBit_dbw_mask < threshold] = 0

    return patternBit_dbw, patternBit_dbw_mask


def test_pattern_loss_mask_mse_2():
    # 示例数据
    bit_num = 1
    # 1.背投影法(BP)计算码阵
    phaseRad_1 = point_2_phase(30, 0)
    phaseRad_2 = point_2_phase(30, 80)
    phaseRad_3 = point_2_phase(30, 150)
    phaseRad_4 = point_2_phase(30, 300)
    # 2.
    threshold = -3
    # patternBit_dbw_1, patternBit_dbw_1_mask = get_mask_by_threshold(bit_num, threshold, phaseRad_1)
    # patternBit_dbw_2, patternBit_dbw_2_mask = get_mask_by_threshold(bit_num, threshold, phaseRad_2)
    # patternBit_dbw_3, patternBit_dbw_3_mask = get_mask_by_threshold(bit_num, threshold, phaseRad_3)
    # patternBit_dbw_4, patternBit_dbw_4_mask = get_mask_by_threshold(bit_num, threshold, phaseRad_4)
    patternBit_dbw_1, patternBit_dbw_1_mask = get_mask_by_slope(bit_num, phaseRad_1)
    patternBit_dbw_2, patternBit_dbw_2_mask = get_mask_by_slope(bit_num, phaseRad_2)
    patternBit_dbw_3, patternBit_dbw_3_mask = get_mask_by_slope(bit_num, phaseRad_3)
    patternBit_dbw_4, patternBit_dbw_4_mask = get_mask_by_slope(bit_num, phaseRad_4)
    print(f"mask-1: count(1)={(patternBit_dbw_1_mask == 1).sum()}, count(0)={(patternBit_dbw_1_mask == 0).sum()}")
    print(f"mask-2: count(1)={(patternBit_dbw_2_mask == 1).sum()}, count(0)={(patternBit_dbw_2_mask == 0).sum()}")
    print(f"mask-3: count(1)={(patternBit_dbw_3_mask == 1).sum()}, count(0)={(patternBit_dbw_3_mask == 0).sum()}")
    print(f"mask-4: count(1)={(patternBit_dbw_4_mask == 1).sum()}, count(0)={(patternBit_dbw_4_mask == 0).sum()}")
    # draw_img(patternBit_dbw_1)
    # draw_img(patternBit_dbw_1_mask)
    # draw_img(patternBit_dbw_2)
    # draw_img(patternBit_dbw_2_mask)
    # draw_img(patternBit_dbw_3)
    # draw_img(patternBit_dbw_3_mask)
    # draw_img(patternBit_dbw_4)
    # draw_img(patternBit_dbw_4_mask)
    # 3.
    phase_mix_12, phaseBit_mix_12, phaseBitDeg_mix_12 \
        = psm_complex_beam_n(phases=[phaseRad_1, phaseRad_2], bit_num=bit_num)
    patternBit_dbw_12 = get_patterndbw(phase_mix_12, bit_num)
    phase_mix_13, phaseBit_mix_13, phaseBitDeg_mix_13 \
        = psm_complex_beam_n(phases=[phaseRad_1, phaseRad_3], bit_num=bit_num)
    patternBit_dbw_13 = get_patterndbw(phase_mix_13, bit_num)
    phase_mix_23, phaseBit_mix_23, phaseBitDeg_mix_23 \
        = psm_complex_beam_n(phases=[phaseRad_2, phaseRad_3], bit_num=bit_num)
    patternBit_dbw_23 = get_patterndbw(phase_mix_23, bit_num)
    phase_mix_34, phaseBit_mix_34, phaseBitDeg_mix_34 \
        = psm_complex_beam_n(phases=[phaseRad_3, phaseRad_4], bit_num=bit_num)
    patternBit_dbw_34 = get_patterndbw(phase_mix_34, bit_num)
    #
    # 计算损失
    # alpha, beta = 100000, 1
    # loss_case_1_good_2 = loss_mask_mse(patternBit_dbw_12,
    #                                    patternBit_dbw_1, patternBit_dbw_2,
    #                                    patternBit_dbw_1_mask, patternBit_dbw_2_mask,
    #                                    alpha, beta)
    # print(f"loss_case_1_good_2: {loss_case_1_good_2}")
    # loss_case_2_good_1_bad_1_13 = loss_mask_mse(patternBit_dbw_13,
    #                                            patternBit_dbw_1, patternBit_dbw_2,
    #                                            patternBit_dbw_1_mask, patternBit_dbw_2_mask,
    #                                            alpha, beta)
    # print(f"loss_case_2_good_1_bad_1_13: {loss_case_2_good_1_bad_1_13}")
    # loss_case_3_good_1_bad_1_23 = loss_mask_mse(patternBit_dbw_23,
    #                                             patternBit_dbw_1, patternBit_dbw_2,
    #                                             patternBit_dbw_1_mask, patternBit_dbw_2_mask,
    #                                             alpha, beta)
    # print(f"loss_case_3_good_1_bad_1_23: {loss_case_3_good_1_bad_1_23}")
    # loss_case_4_bad_2 = loss_mask_mse(patternBit_dbw_34,
    #                                   patternBit_dbw_1, patternBit_dbw_2,
    #                                   patternBit_dbw_1_mask, patternBit_dbw_2_mask,
    #                                   alpha, beta)
    # print(f"loss_case_4_bad_2: {loss_case_4_bad_2}")
    #
    alpha, beta = 100000, 1
    loss_case_1_good_2 = loss_mask_mse_2(patternBit_dbw_12, patternBit_dbw_1_mask, patternBit_dbw_2_mask, alpha, beta)
    print(f"loss_case_1_good_2: {loss_case_1_good_2}")
    loss_case_2_good_1_bad_1_13 = loss_mask_mse_2(patternBit_dbw_13, patternBit_dbw_1_mask, patternBit_dbw_2_mask, alpha, beta)
    print(f"loss_case_2_good_1_bad_1_13: {loss_case_2_good_1_bad_1_13}")
    loss_case_3_good_1_bad_1_23 = loss_mask_mse_2(patternBit_dbw_23, patternBit_dbw_1_mask, patternBit_dbw_2_mask, alpha, beta)
    print(f"loss_case_3_good_1_bad_1_23: {loss_case_3_good_1_bad_1_23}")
    loss_case_4_bad_2 = loss_mask_mse_2(patternBit_dbw_34, patternBit_dbw_1_mask, patternBit_dbw_2_mask, alpha, beta)
    print(f"loss_case_4_bad_2: {loss_case_4_bad_2}")
    #
    # alpha, beta = 50000, 1
    # loss_case_1_good_2 = loss_mask_ssim(patternBit_dbw_12,
    #                                    patternBit_dbw_1, patternBit_dbw_2,
    #                                    patternBit_dbw_1_mask, patternBit_dbw_2_mask,
    #                                    alpha, beta)
    # print(f"loss_case_1_good_2: {loss_case_1_good_2}")
    # loss_case_2_good_1_bad_1_13 = loss_mask_ssim(patternBit_dbw_13,
    #                                             patternBit_dbw_1, patternBit_dbw_2,
    #                                             patternBit_dbw_1_mask, patternBit_dbw_2_mask,
    #                                             alpha, beta)
    # print(f"loss_case_2_good_1_bad_1_13: {loss_case_2_good_1_bad_1_13}")
    # loss_case_3_good_1_bad_1_23 = loss_mask_ssim(patternBit_dbw_23,
    #                                             patternBit_dbw_1, patternBit_dbw_2,
    #                                             patternBit_dbw_1_mask, patternBit_dbw_2_mask,
    #                                             alpha, beta)
    # print(f"loss_case_3_good_1_bad_1_23: {loss_case_3_good_1_bad_1_23}")
    # loss_case_4_bad_2 = loss_mask_ssim(patternBit_dbw_34,
    #                                   patternBit_dbw_1, patternBit_dbw_2,
    #                                   patternBit_dbw_1_mask, patternBit_dbw_2_mask,
    #                                   alpha, beta)
    # print(f"loss_case_4_bad_2: {loss_case_4_bad_2}")



def main_loss_adapter_mse_ssim_ncc_item(bit_num, theta, phi):
    # 1.背投影法(BP)计算码阵
    phaseRad = point_2_phase(theta, phi)
    patternBit_dbw = get_patterndbw(phaseRad, bit_num)
    # pattern 转dB
    patternBit_dbw_threld_30 = np.copy(patternBit_dbw)
    patternBit_dbw_threld_100 = np.copy(patternBit_dbw)
    # 将 patternBit_dbw 中所有小于 -3 的值设置为最小值
    patternBit_dbw_threld_30[patternBit_dbw_threld_30 < -3] = -30
    patternBit_dbw_threld_100[patternBit_dbw_threld_100 < -3] = -100
    return phaseRad, patternBit_dbw, patternBit_dbw_threld_30, patternBit_dbw_threld_100


def main_loss_adapter_mse_ssim_ncc(logger):
    bit_num = 1
    phaseRad1, patternBit1_dbw, patternBit1_dbw_threld_30, patternBit1_dbw_threld_100 \
        = main_loss_adapter_mse_ssim_ncc_item(bit_num, 45, 270)
    phaseRad2, patternBit2_dbw, patternBit2_dbw_threld_30, patternBit2_dbw_threld_100 \
        = main_loss_adapter_mse_ssim_ncc_item(bit_num, 30, 60)
    # 画图检查
    # draw_img(phaseRad1)
    # draw_img(patternBit1_dbw)
    # draw_img(patternBit1_dbw_threld_30)
    # draw_img(patternBit1_dbw_threld_100)
    # patterndbw1, x1, y1 = get_patterndbw_xy(phaseRad1)
    # draw_img_xyz_2(patterndbw1, x1, y1)
    # draw_img(phaseRad2)
    # draw_img(patternBit2_dbw)
    # draw_img(patternBit2_dbw_threld_30)
    # draw_img(patternBit2_dbw_threld_100)
    # patterndbw2, x2, y2 = get_patterndbw_xy(phaseRad2)
    # draw_img_xyz_2(patterndbw2, x2, y2)
    # 算loss
    logger.info("1 vs 1:")
    loss_adapter_mse_ssim_ncc(patternBit1_dbw_threld_100, phaseRad1)
    loss_adapter_mse_ssim_ncc(patternBit1_dbw_threld_30, phaseRad1)
    loss_adapter_mse_ssim_ncc(patternBit1_dbw, phaseRad1)
    logger.info("2 vs 2:")
    loss_adapter_mse_ssim_ncc(patternBit2_dbw_threld_100, phaseRad2)
    loss_adapter_mse_ssim_ncc(patternBit2_dbw_threld_30, phaseRad2)
    loss_adapter_mse_ssim_ncc(patternBit2_dbw, phaseRad2)
    logger.info("1 vs 2:")
    loss_adapter_mse_ssim_ncc(patternBit1_dbw_threld_100, phaseRad2)
    loss_adapter_mse_ssim_ncc(patternBit1_dbw_threld_30, phaseRad2)
    loss_adapter_mse_ssim_ncc(patternBit1_dbw, phaseRad2)


if __name__=="__main__":
    # 初始化日志
    logger = setup_logging()
    #
    # main_loss_adapter_mse_ssim_ncc(logger)
    # main_ga(logger)
    #
    # test_loss_mask_mse()
    # test_pattern_loss_mask_mse()
    test_pattern_loss_mask_mse_2()


